﻿using Microsoft.Extensions.AI;
namespace Qwen3;

class Qwen3MoEChatClient
{
    private readonly Qwen3ChatTemplateBuilder _chatTemplateBuilder;
    private readonly Qwen3MoEPipeline _pipeline;

    public Qwen3MoEChatClient(Qwen3MoEPipeline pipeline)
    {
        _chatTemplateBuilder = Qwen3ChatTemplateBuilder.Instance;
        _pipeline = pipeline;
    }

    public Task<ChatResponse> GetResponseAsync(IEnumerable<ChatMessage> messages, ChatOptions? options = null, CancellationToken cancellationToken = default)
    {
        var prompt = _chatTemplateBuilder.BuildPrompt(messages, options);
        var stopTokens = new List<int>();
        stopTokens.Add(151645);
        stopTokens.Add(151644);
        stopTokens.Add(151643);
        var output = _pipeline.Generate(prompt,
            stopTokens: stopTokens,
            maxLen: options?.MaxOutputTokens ?? 1024,
            temperature: options?.Temperature ?? 0.7f,
            topP: options?.TopP ?? 0.9f);
        var chatMessage = new ChatMessage(ChatRole.Assistant, output);
        return Task.FromResult(new ChatResponse([chatMessage])
        {
            CreatedAt = DateTime.UtcNow,
            FinishReason = ChatFinishReason.Stop,
            ResponseId = Guid.NewGuid().ToString("N"),
        });
    }

    public async IAsyncEnumerable<ChatResponseUpdate> GetStreamingResponseAsync(
        IEnumerable<ChatMessage> messages,
        ChatOptions? options = null)
    {
        var prompt = _chatTemplateBuilder.BuildPrompt(messages, options);
        var stopTokens = new List<int>();
        stopTokens.Add(151645);
        stopTokens.Add(151644);
        stopTokens.Add(151643);
        string responseId = Guid.NewGuid().ToString("N");
        foreach (var output in  _pipeline.GenerateStreaming(prompt,
            stopTokens: stopTokens,
            maxLen: options?.MaxOutputTokens ?? 1024,
            temperature: options?.Temperature ?? 0.7f,
            topP: options?.TopP ?? 0.9f))
        {
            yield return new(ChatRole.Assistant, output)
            {
                CreatedAt = DateTime.UtcNow,
                ResponseId = responseId,
            };
        }
    }
}